package com.github.yoojia.mqtt;

import com.github.yoojia.mqtt.util.MapCached;
import com.github.yoojia.mqtt.util.SnowflakeId;
import com.parkingwang.lang.*;
import com.parkingwang.lang.data.Pair;
import com.parkingwang.lang.kit.StringKit;
import org.eclipse.paho.client.mqttv3.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

import static java.util.Objects.requireNonNull;

/**
 * @author Yoojia Chen (yoojiachen@gmail.com)
 * @since 1.0.0
 */
public class MQTTSocket {

    private static final Logger LOGGER = LoggerFactory.getLogger(MQTTSocket.class);

    public static final String TOPIC_TYPE_REQ = "requests";
    public static final String TOPIC_TYPE_REP = "replies";
    public static final String TOPIC_TYPE_PUB_SUB = "pubsub";

    private final SnowflakeId mSnowflakeId;
    private final Config mConfig;

    private final String mTopicREQ;
    private final String mTopicREP;
    private final String mTopicPUBSUB;

    private static final SocketListener NOP_SOCKET_LISTENER = new SocketListener() {
        @Override
        public void onConnected() {
        }

        @Override
        public void onConnectError(Throwable error) {
        }

        @Override
        public void onConnectionLost(Throwable cause) {
        }

        @Override
        public void onDisconnected() {
        }

        @Override
        public void onMessageError(IOException e) {
        }
    };

    private final LazyArg<MqttClient, Pair<String, String>> mClient = LazyArg.from(new SupplierArg<MqttClient, Pair<String, String>>() {
        @Override
        public MqttClient call(Pair<String, String> args) {
            final String address = args.first;
            final String clientId = args.second;
            return Try.die(new SupplierThrow<MqttClient>() {
                @Override
                public MqttClient call() throws Throwable {
                    return new MqttClient(address, clientId);
                }
            });
        }
    });

    /**
     * 懒加载的消息订阅器。
     * 注意：此订阅器为懒加载对象。当首次调用时才初始化并建立订阅通道。
     */
    private final MapCached<String, Subscriber> mLazySubscribers = new MapCached<String, Subscriber>(3) {
        @Override
        protected Subscriber newByKey(String topicFilter) {
            // 初始化订阅器时，发起订阅请求。
            final Subscriber sub = new Subscriber(MQTTSocket.this, mSnowflakeId, mClient.getPresent(), topicFilter);
            LOGGER.info("Lazy initial subscriber... Call to subscribe topic: {}", topicFilter);
            try {
                sub.subscribe();
            } catch (MqttException e) {
                LOGGER.error("Fail to subscribe topic: " + topicFilter, e);
                throw new IllegalStateException("SUBSCRIBE FAILED", e);
            }
            return sub;
        }
    };

    MQTTSocket(Config config) {
        mConfig = config;
        mSnowflakeId = new SnowflakeId(mConfig.workerId, mConfig.dataCenterId);
        // 订阅主题基于三个元素：Domain, NodeId, MessageType
        mTopicREQ = TopicsOf.subscribe(mConfig.domain, mConfig.nodeId, TOPIC_TYPE_REQ);
        mTopicREP = TopicsOf.subscribe(mConfig.domain, mConfig.nodeId, TOPIC_TYPE_REP);
        mTopicPUBSUB = TopicsOf.subscribe(mConfig.domain, mConfig.nodeId, TOPIC_TYPE_PUB_SUB);
    }

    public Config getConfig() {
        return mConfig;
    }

    public MqttClient getMqttClient() {
        return mClient.getPresent();
    }

    //// 连接服务端接口：

    public void connect() {
        final MqttConnectOptions opts = new MqttConnectOptions();
        opts.setAutomaticReconnect(mConfig.autoReconnect);
        opts.setCleanSession(mConfig.clearSession);
        opts.setConnectionTimeout(mConfig.timeout <= 0 ? 60 : mConfig.timeout);
        if (StringKit.isNotEmpty(mConfig.username)) {
            opts.setUserName(mConfig.username);
        }
        if (StringKit.isNotEmpty(mConfig.password)) {
            opts.setPassword(mConfig.password.toCharArray());
        }
        final MqttClient client = mClient.getCheckedWithArg(new Pair<>(mConfig.address, mConfig.nodeId));
        if (StringKit.isNotEmpty(mConfig.username) && StringKit.isNotEmpty(mConfig.password)) {
            LOGGER.info("Connecting to broker, Address= {}, nodeId= {}, Auth.username= {}, Auth.password= {}",
                    mConfig.address, mConfig.nodeId, mConfig.username, mConfig.password);
        } else {
            LOGGER.info("Connecting to broker, Address= {}, nodeId= {}",
                    mConfig.address, mConfig.nodeId);
        }
        // 监听连接状态
        client.setCallback(new MqttCallback() {
            @Override
            public void connectionLost(Throwable cause) {
                LOGGER.error("Broker connection lost", cause);
                getSocketListener().onConnectionLost(cause);
            }

            @Override
            public void messageArrived(String topic, MqttMessage message) throws Exception {
            }

            @Override
            public void deliveryComplete(IMqttDeliveryToken token) {
            }
        });
        // 重试3次连接
        Try.retry(3, new Try.RetryAction() {
            @Override
            public void onWork() throws Throwable {
                client.connect(opts);
                // 连接成功后，添加自动断开连接的处理
                getSocketListener().onConnected();
                Runtime.getRuntime()
                        .addShutdownHook(new Thread(new Runnable() {
                            @Override
                            public void run() {
                                if (mClient.isPresent()) {
                                    unsubscribeTopicsAndDisconnect(mClient.getPresent());
                                }
                            }
                        }));
            }

            @Override
            public void onThrows(Throwable err, int retry) {
                getSocketListener().onConnectError(err);
                LOGGER.error("Error when connect to broker, retry: " + retry, err);
            }

            @Override
            public void onFailed() {
                LOGGER.error("Fail to connect, retry...");
            }
        }, 1000);
    }

    public void disconnect() {
        mClient.ifPresent(new Consumer<MqttClient>() {
            @Override
            public void call(MqttClient client) {
                unsubscribeTopicsAndDisconnect(client);
            }
        });
    }

    //// 发送消息：

    /**
     * 发送Request-Reply消息请求
     *
     * @param request Request消息
     * @return Call接口，可以按需要同步或异步操作
     */
    public Call sendRequestCall(Message request) {
        return new RealCall(mLazySubscribers.getOrNew(mTopicREP),
                checkMessage(request),
                this);
    }

    /**
     * 发送消息，如果发送消息发生错误，则抛出IOException异常
     *
     * @param message 消息
     * @throws IOException 发送消息发生错误时抛出
     */
    public void sendThrows(Message message) throws IOException {
        final Message publish = checkMessage(message);
        final MqttMessage realMsg = new MqttMessage(publish.payload);
        realMsg.setQos(publish.qos);
        realMsg.setRetained(publish.retained);
        final String topic = TopicsOf.publish(message.domain, message.targetNodeId, message.type,
                message.messageId, message.senderNodeId, message.tag);
        try {
            LOGGER.debug("Publish MQTT message, topic={}, msg={}", topic, realMsg);
            mClient.getPresent().publish(topic, realMsg);
        } catch (MqttException e) {
            final IOException error = new IOException(e);
            getSocketListener().onMessageError(error);
            throw error;
        }
    }

    /**
     * 发送消息
     *
     * @param message 消息对象
     */
    public void send(Message message) {
        try {
            sendThrows(message);
        } catch (IOException e) {/*在sendThrows中已处理*/}
    }

    //// 消息处理器注册接口：

    /**
     * 注册一个指定Topic的请求消息处理器，指定本地消息过滤器和消息处理器。
     *
     * @param topic   MQTT的订阅Topic
     * @param filter  本地消息过滤器
     * @param handler 消息处理器
     * @return 消息处理器的流水编号，用于后续反注册消息处理器
     */
    public long addMessageHandler(String topic, LocalFilter filter, MessageHandler handler) {
        requireNonNull(topic);
        requireNonNull(filter);
        requireNonNull(handler);
        return mLazySubscribers.getOrNew(topic)
                .addHandler(filter, handler);
    }

    /**
     * 取消指定MQTT订阅Topic及注册流水编号的消息处理器
     *
     * @param topic MQTT的订阅Topic
     * @param regId 消息处理流水编号
     */
    public void removeMessageHandler(String topic, final long regId) {
        mLazySubscribers.ifPresent(topic, new Consumer<Subscriber>() {
            @Override
            public void call(Subscriber subscriber) {
                subscriber.removeHandler(regId);
            }
        });
    }

    /**
     * 注册一个Request请求消息处理器，指定本地消息过滤器和消息处理器。
     *
     * @param filter  本地消息过滤器
     * @param handler 消息处理器
     * @return 消息处理器的流水编号，用于后续反注册消息处理器
     */
    public long addRequestMessageHandler(LocalFilter filter, MessageHandler handler) {
        return addMessageHandler(mTopicREQ, filter, handler);
    }

    /**
     * 注册一个Request请求消息处理器，指定本地消息过滤器和消息处理器。
     *
     * @param handler 消息处理器
     * @return 消息处理器的流水编号，用于后续反注册消息处理器
     */
    public long addRequestMessageHandler(MessageHandler handler) {
        return addRequestMessageHandler(LocalFilter.ANY, handler);
    }

    /**
     * 取消Request的消息处理器
     *
     * @param regId 消息处理流水编号
     */
    public void removeRequestMessageHandler(long regId) {
        removeMessageHandler(mTopicREQ, regId);
    }

    /**
     * 注册一个Reply请求消息处理器，指定本地消息过滤器和消息处理器。
     *
     * @param filter  本地消息过滤器
     * @param handler 消息处理器
     * @return 消息处理器的流水编号，用于后续反注册消息处理器
     */
    public long addReplyMessageHandler(LocalFilter filter, MessageHandler handler) {
        return addMessageHandler(mTopicREP, filter, handler);
    }

    /**
     * 注册一个Reply请求消息处理器。
     *
     * @param handler 消息处理器
     * @return 消息处理器的流水编号，用于后续反注册消息处理器
     */
    public long addReplyMessageHandler(MessageHandler handler) {
        return addReplyMessageHandler(LocalFilter.ANY, handler);
    }

    /**
     * 取消Reply的消息处理器
     *
     * @param regId 消息处理流水编号
     */
    public void removeReplyMessageHandler(long regId) {
        removeMessageHandler(mTopicREP, regId);
    }

    /**
     * 注册一个Event请求消息处理器，指定本地消息过滤器和消息处理器。
     *
     * @param filter  本地消息过滤器
     * @param handler 消息处理器
     * @return 消息处理器的流水编号，用于后续反注册消息处理器
     */
    public long addSubMessageHandler(LocalFilter filter, MessageHandler handler) {
        return addMessageHandler(mTopicPUBSUB, filter, handler);
    }

    /**
     * 注册一个Event请求消息处理器。
     *
     * @param handler 消息处理器
     * @return 消息处理器的流水编号，用于后续反注册消息处理器
     */
    public long addSubMessageHandler(MessageHandler handler) {
        return addSubMessageHandler(LocalFilter.ANY, handler);
    }

    /**
     * 取消Event的消息处理器
     *
     * @param regId 消息处理流水编号
     */
    public void removeSubMessageHandler(long regId) {
        removeMessageHandler(mTopicPUBSUB, regId);
    }

    //// 构建消息对象：

    /**
     * 获取下一个Request消息ID
     *
     * @return Request消息ID
     */
    public long nextRequestId() {
        return mSnowflakeId.nextId();
    }

    /**
     * 获取一个发送给指定节点ID的Request类型消息对象，并指定Payload。
     *
     * @param targetNodeId 目标节点ID
     * @param payload      消息的负载
     * @return Message对象
     */
    public Message newRequestMessageFor(String targetNodeId, byte[] payload) {
        return newMessage(targetNodeId, TOPIC_TYPE_REQ, payload)
                .builder()
                .messageId(nextRequestId()) // 填充Request的请求ID
                .build();
    }

    /**
     * 基于Request消息，获取它的Reply消息
     *
     * @param request Request消息
     * @param payload 响应消息负载
     * @return Message对象
     */
    public Message newReplyMessageOf(Message request, byte[] payload) {
        requireNonNull(request);
        return new Message.Builder()
                .domain(mConfig.domain)
                .targetNodeId(request.senderNodeId) // 返回响应的客户端ID为消息的来源终端
                .type(TOPIC_TYPE_REP)
                .messageId(request.messageId) // 返回响应消息的ID保持不变
                .senderNodeId(mConfig.nodeId) // 发送消息的终端为本连接Channel
                .payload(payload)
                .qos(request.qos)
                .retained(request.retained)
                .build();
    }

    /**
     * 获取一个PUB_SUB类型的消息
     *
     * @param targetNodeId 目标NodeId
     * @param payload      消息的负载
     * @return Message消息对象
     */
    public Message newPubMessageFor(String targetNodeId, byte[] payload) {
        return newMessage(targetNodeId, TOPIC_TYPE_PUB_SUB, payload);
    }

    /**
     * 获取一个本客户端的PUB_SUB类型消息
     *
     * @param payload 消息负载
     * @return Message消息对象
     */
    public Message newPubMessageOfSelf(byte[] payload) {
        return newPubMessageFor(mConfig.nodeId, payload);
    }

    /**
     * 获取一个消息对象，指定目标节点ID、消息类型及负载
     *
     * @param targetNodeId 目标节点ID
     * @param type         消息类型
     * @param payload      负载
     * @return Message消息对象
     */
    public Message newMessage(String targetNodeId, String type, byte[] payload) {
        requireNonNull(targetNodeId);
        requireNonNull(type);
        requireNonNull(payload);
        return new Message.Builder()
                .domain(mConfig.domain)
                .targetNodeId(targetNodeId) // 指定目标Channel
                .type(type)
                .messageId(0)
                .senderNodeId(mConfig.nodeId) // 发送消息的终端为本连接
                .payload(payload)
                .build();
    }

    //////

    private void unsubscribeTopicsAndDisconnect(final MqttClient client) {
        if (client.isConnected()) {
            // Unsubscribe All
            if (mConfig.unsubscribeOnDisconnect) {
                for (final Subscriber sub : mLazySubscribers.values()) {
                    Try.ignore(new ActionThrow.Action() {
                        @Override
                        public void call() throws Throwable {
                            client.unsubscribe(sub.getTopicFilter());
                        }
                    });
                }
                LOGGER.info("Unsubscribe all topic");
            }
            Try.ignore(new ActionThrow.Action() {
                @Override
                public void call() throws Throwable {
                    client.disconnect();
                    getSocketListener().onDisconnected();
                }
            });
            LOGGER.info("Disconnect from broker");
        }
    }

    private SocketListener getSocketListener() {
        if (mConfig.socketListener != null) {
            return mConfig.socketListener;
        } else {
            return NOP_SOCKET_LISTENER;
        }
    }

    private static Message checkMessage(Message message) {
        requireNonNull(message);
        if (StringKit.isEmpty(message.domain)) throw new IllegalArgumentException("Message Field 'domain' is required");
        if (StringKit.isEmpty(message.targetNodeId))
            throw new IllegalArgumentException("Message Field 'targetNodeId' is required");
        if (StringKit.isEmpty(message.type)) throw new IllegalArgumentException("Message Field 'type' is required");
        if (message.messageId < 0) throw new IllegalArgumentException("Message Field 'messageId' must be >= 0");
        if (StringKit.isEmpty(message.senderNodeId))
            throw new IllegalArgumentException("Message Field 'senderNodeId' is required");
        return message;
    }

    public static Context context() {
        return new Context();
    }

    //////

    static class Config {

        public final String domain;
        public final String nodeId;
        public final String username;
        public final String password;
        public final String address;
        public final int timeout;
        public final boolean autoReconnect;
        public final boolean clearSession;
        public final int dataCenterId;
        public final int workerId;
        public final boolean unsubscribeOnDisconnect;
        public final SocketListener socketListener;

        Config(String domain, String nodeId, String username, String password, String address,
               int timeout, boolean autoReconnect, boolean clearSession,
               int dataCenterId, int workerId, boolean unsubscribeOnDisconnect,
               SocketListener socketListener) {
            this.domain = domain;
            this.nodeId = nodeId;
            this.username = username;
            this.password = password;
            this.address = address;
            this.timeout = timeout;
            this.autoReconnect = autoReconnect;
            this.clearSession = clearSession;
            this.dataCenterId = dataCenterId;
            this.workerId = workerId;
            this.unsubscribeOnDisconnect = unsubscribeOnDisconnect;
            this.socketListener = socketListener;
        }
    }
}
